package storage.impl;

import org.wwx.register.model.ServiceInstance;
import storage.ServiceStorage;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.locks.StampedLock;

/**
 * 修复并发问题的高性能内存存储实现
 *
 * @date 2025/6/6
 * @author王闻薪
 */
public class InMemoryStorage implements ServiceStorage {

    // 使用嵌套的 ConcurrentHashMap 确保服务级别的隔离
    private final Map<String, ServiceEntry> serviceRegistry = new ConcurrentHashMap<>();

    // 实例ID到服务名的映射，用于快速定位
    private final Map<String, String> instanceToServiceMap = new ConcurrentHashMap<>();

    // 使用 StampedLock 提供高性能的读写锁
    private final StampedLock globalLock = new StampedLock();

    // 服务级别的内部存储结构
    private static class ServiceEntry {
        // 使用 CopyOnWriteArrayList 保证读操作的并发性能
        private final List<ServiceInstance> instances = new CopyOnWriteArrayList<>();

        // 服务级别的锁，减少全局锁竞争
        private final Object serviceLock = new Object();

        // 实例ID到实例的映射，用于快速查找
        private final Map<String, ServiceInstance> instanceMap = new HashMap<>();
    }

    @Override
    public void registerService(String serviceName, ServiceInstance instance) {
        long stamp = globalLock.writeLock();
        try {
            // 双重检查确保服务条目存在
            ServiceEntry entry = serviceRegistry.computeIfAbsent(
                    serviceName,
                    k -> new ServiceEntry()
            );

            synchronized (entry.serviceLock) {
                // 检查实例是否已存在
                if (entry.instanceMap.containsKey(instance.getId())) {
                    throw new IllegalStateException("Instance already registered: " + instance.getId());
                }

                // 添加实例到服务条目
                entry.instances.add(instance);
                entry.instanceMap.put(instance.getId(), instance);

                // 更新全局实例映射
                instanceToServiceMap.put(instance.getId(), serviceName);
            }

            // 确保内存修改对所有线程可见
            synchronized (this) {
                // 内存屏障
            }
        } finally {
            globalLock.unlockWrite(stamp);
        }
    }

    @Override
    public void unregisterService(String serviceName, String instanceId) {
        long stamp = globalLock.writeLock();
        try {
            ServiceEntry entry = serviceRegistry.get(serviceName);
            if (entry == null) return;

            synchronized (entry.serviceLock) {
                ServiceInstance instance = entry.instanceMap.get(instanceId);
                if (instance != null) {
                    // 原子性移除操作
                    entry.instances.removeIf(i -> i.getId().equals(instanceId));
                    entry.instanceMap.remove(instanceId);
                    instanceToServiceMap.remove(instanceId);

                    // 如果服务没有实例，清理服务条目
                    if (entry.instances.isEmpty()) {
                        serviceRegistry.remove(serviceName);
                    }
                }
            }

            // 内存屏障确保可见性
            synchronized (this) {
                // 空同步块仅用于内存屏障
            }
        } finally {
            globalLock.unlockWrite(stamp);
        }
    }

    @Override
    public List<ServiceInstance> getServiceInstances(String serviceName) {
        // 使用乐观读锁提高性能
        long stamp = globalLock.tryOptimisticRead();
        ServiceEntry entry = serviceRegistry.get(serviceName);

        if (!globalLock.validate(stamp)) {
            // 乐观读失败，升级为悲观读锁
            stamp = globalLock.readLock();
            try {
                entry = serviceRegistry.get(serviceName);
            } finally {
                globalLock.unlockRead(stamp);
            }
        }

        if (entry == null) {
            return Collections.emptyList();
        }

        // 返回实例的不可修改视图
        return Collections.unmodifiableList(new ArrayList<>(entry.instances));
    }

    @Override
    public void renewLease(String serviceName, String instanceId) {
        // 使用服务名和实例ID双重验证
        String actualServiceName = instanceToServiceMap.get(instanceId);
        if (actualServiceName == null || !actualServiceName.equals(serviceName)) {
            return;
        }

        ServiceEntry entry = serviceRegistry.get(serviceName);
        if (entry == null) return;

        // 使用读锁保护
        long stamp = globalLock.readLock();
        try {
            ServiceInstance instance = entry.instanceMap.get(instanceId);
            if (instance != null) {
                // 使用 volatile 语义更新
                synchronized (instance) {
                    instance.renewLease();
                }
            }
        } finally {
            globalLock.unlockRead(stamp);
        }
    }

    // 辅助方法：获取所有服务名称
    public Set<String> getAllServiceNames() {
        long stamp = globalLock.readLock();
        try {
            return new HashSet<>(serviceRegistry.keySet());
        } finally {
            globalLock.unlockRead(stamp);
        }
    }

    // 辅助方法：获取所有实例
    public List<ServiceInstance> getAllInstances() {
        long stamp = globalLock.readLock();
        try {
            List<ServiceInstance> allInstances = new ArrayList<>();
            for (ServiceEntry entry : serviceRegistry.values()) {
                allInstances.addAll(entry.instances);
            }
            return Collections.unmodifiableList(allInstances);
        } finally {
            globalLock.unlockRead(stamp);
        }
    }

    // 确保内存可见性的工具方法
    private void ensureVisibility() {
        synchronized (this) {
            // 空同步块仅用于内存屏障
        }
    }
}